import copy
import bson

from pymongo import MongoClient, UpdateOne

from utility.log import logger
from conf.settings import CLOUDMONITOR_MONGODB, WASS_MONGODB, WASS_AUTHENTICATE

TARGET_DB_DICT = {
    'host_port': {'host': WASS_MONGODB},
    'auth': WASS_AUTHENTICATE,
    'db_name': 'wass',
    'collection_tasks': 'scan_task_sessions',
    'collection_event': 'scan_events',
    'collection_task': 'scan_tasks',
}

LOCAL_DB_DICT = {
    'host_port': {'host': CLOUDMONITOR_MONGODB},
    'db_name': 'cloudmonitor',
    'collection_tasks': 'tasks'
}

# scan_task_sessions表对应task表的字段
KEY_MAP = {
    'status': 'status',
    'task_session_id': 'taskId',
    'progress': 'progress',
    'created_at': 'startTime',
    'ended_at': 'endTime',
}

# scan_event表对应task表的字段
EVENT_MAP = {
    # 'target_url': 'target',
}


# scan_task表对应task表字段
SCAN_TASK_MAP = {
    "next_start_at": "nextExecDate"
}


VALUE_MAP = {
    'status': {
        3: 'active',
        9: 'completed',
        10: 'failed'
    }
}


class Task(object):
    addition_uid_list = ['uid', 'refId', 'jobId', 'prevTaskId']

    def __init__(self, collection):
        self.collection = collection
        self.task_list = []
        self.other_update_list = []

    def update_task(self, data_list) -> (list, list):
        session_id_list = []
        task_id_list = []

        for item in data_list:
            addition = item.get('addition', {})
            addition.pop('is_sync')
            _id = addition.get('_id')
            session_id_list.append(item.get('task_session_id'))
            task_id_list.append(item.get('task_id'))

            if _id:
                _id = bson.ObjectId(_id)

            temp_dict = dict()
            for k in KEY_MAP.keys():
                v = item.get(k, "")
                temp_dict[KEY_MAP[k]] = VALUE_MAP.get(k, {}).get(v, v)

            for k, v in addition.items():
                if k in self.addition_uid_list:
                    addition[k] = bson.ObjectId(v)

            temp_dict.update(addition)

            # set default value
            temp_dict['result'] = dict()
            temp_dict['securityStatus'] = 'unknown'
            temp_dict['isLatestTask'] = self._is_latest_task(item, temp_dict)

            self.task_list.append((item.get('task_id'), temp_dict))

        return session_id_list, task_id_list

    def update_results(self, results_dict, task_results_list):
        for _task_id, _item in self.task_list:
            _session_id = _item.get(KEY_MAP.get('task_session_id'))
            _item.update(results_dict.get(_session_id, {}))
            _item.update(task_results_list.get(_task_id, {}))

    def save(self):
        update_list = []
        for _task_id, _task in self.task_list:
            _session_id = _task.get(KEY_MAP.get('task_session_id'))
            if _session_id:
                update_list.append(UpdateOne({'taskId': _session_id}, {'$set': _task}, upsert=True))

        if self.other_update_list:
            self.collection.bulk_write(self.other_update_list)
        if update_list:
            self.collection.bulk_write(update_list)

    def _is_latest_task(self, session_data, task_data):
        if session_data.get('status') == 9:
            condition = {
                'target': task_data.get('target'),
                'uid': task_data.get('uid'),
                'isLatestTask': True
            }
            self.other_update_list.append(UpdateOne(condition, {'$set': {'isLatestTask': False}}))
            return True

        return False

    # @staticmethod
    # def _get_task_type(item):
    #     task_type = item.get('flex_job_config_map', {}).get('cloud', {}).get('probe_type')
    #     if task_type:
    #         return task_type.lower()
    #
    #     task_type = item.get('flex_job_config_map', {}).get('appscan')
    #     if task_type:
    #         return 'securityEvent'


class ScanTaskSessions(object):
    def __init__(self, collection):
        self.collection = collection
        self.update_list = []

        condition = {
            'status': {'$in': [3, 9, 10]},
            'addition.is_sync': '0'
        }
        self.data = self.collection.find(condition)

    def get_data(self):
        data_list = []
        for _data in self.data:
            _data['addition']['is_sync'] = '1'

            data_list.append(copy.deepcopy(_data))
            self.update_list.append(UpdateOne({'_id': _data['_id']}, {'$set': _data}))

        return data_list

    def save(self):
        if self.update_list:
            self.collection.bulk_write(self.update_list)


class ScanEvents(object):
    def __init__(self, collection):
        self.collection = collection

    def get_results(self, session_id_list) -> dict:
        results = dict()
        for _item in self.collection.find({'task_session_id': {'$in': session_id_list}}):

            detail = _item.get('detail', {})
            if detail.get('cloud'):
                _temp = self._get_detail_cloud(detail)
            else:
                _temp = self._get_detail_appscan(detail)

            for k, v in EVENT_MAP.items():
                _value = _item.get(k)
                if _value is not None:
                    _temp[v] = _value

            results[_item['task_session_id']] = _temp

        return results

    @staticmethod
    def _get_detail_cloud(detail):
        return {
            'result': {
                'addition': {
                    'detail': detail.get('cloud', [])
                }
            }
        }

    @staticmethod
    def _get_detail_appscan(detail):
        return {
            'result': {
                'addition': detail
            }
        }


class ScanTask(object):
    def __init__(self, collection):
        self.collection = collection

    def get_results(self, task_id_list) -> dict:
        results = dict()
        for _item in self.collection.find({'task_id': {'$in': task_id_list}}):
            if _item.get('status') in [9, 10]:
                temp_dict = dict()
                for k, v in SCAN_TASK_MAP.items():
                    temp_dict[v] = _item.get(k)

                results[_item['task_id']] = temp_dict

        return results


def main():
    logger.info(f'Start {__file__}')
    try:
        conn = MongoClient(**LOCAL_DB_DICT['host_port'])
        conn_target = MongoClient(**TARGET_DB_DICT['host_port'], serverSelectionTimeoutMS=61000, socketTimeoutMS=61000)
        conn_target['admin'].authenticate(**TARGET_DB_DICT['auth'])

        task = Task(conn[LOCAL_DB_DICT['db_name']][LOCAL_DB_DICT['collection_tasks']])

        scan_task_sessions = ScanTaskSessions(conn_target[TARGET_DB_DICT['db_name']][TARGET_DB_DICT['collection_tasks']])
        scan_events = ScanEvents(conn_target[TARGET_DB_DICT['db_name']][TARGET_DB_DICT['collection_event']])
        scan_tasks = ScanTask(conn_target[TARGET_DB_DICT['db_name']][TARGET_DB_DICT['collection_task']])

        data_list = scan_task_sessions.get_data()

        session_list, task_id_list = task.update_task(data_list)

        result_list = scan_events.get_results(session_list)
        task_results_list = scan_tasks.get_results(task_id_list)

        task.update_results(result_list, task_results_list)

        task.save()
        scan_task_sessions.save()

        conn.close()
        conn_target.close()

        logger.info(f'Complete {__file__}')
    except Exception as e:
        logger.exception(e)


if __name__ == '__main__':
    main()
